package com.macro.mall.common.util;

import org.slf4j.MDC;
import org.springframework.cloud.sleuth.Span;
import org.springframework.cloud.sleuth.Tracer;

import java.util.ArrayList;
import java.util.Collection;
import java.util.List;
import java.util.Map;
import java.util.concurrent.*;

/**
 * MDC工具类，用于在多线程环境下传递traceId等上下文信息
 * @author macro
 */
public class MDCUtil {

    /**
     * 将MDC上下文信息封装到Runnable中，确保子线程能获取到父线程的MDC上下文
     * @param runnable 原始Runnable
     * @return 包装后的Runnable
     */
    public static Runnable wrap(final Runnable runnable) {
        // 获取父线程的MDC上下文
        Map<String, String> contextMap = MDC.getCopyOfContextMap();
        
        return () -> {
            try {
                // 将父线程的MDC上下文设置到子线程中
                if (contextMap != null) {
                    MDC.setContextMap(contextMap);
                }
                runnable.run();
            } finally {
                // 清理子线程的MDC上下文
                MDC.clear();
            }
        };
    }
    
    /**
     * 将MDC上下文信息封装到Callable中，确保子线程能获取到父线程的MDC上下文
     * @param callable 原始Callable
     * @param <T> 返回值类型
     * @return 包装后的Callable
     */
    public static <T> Callable<T> wrap(final Callable<T> callable) {
        // 获取父线程的MDC上下文
        Map<String, String> contextMap = MDC.getCopyOfContextMap();
        
        return () -> {
            try {
                // 将父线程的MDC上下文设置到子线程中
                if (contextMap != null) {
                    MDC.setContextMap(contextMap);
                }
                return callable.call();
            } finally {
                // 清理子线程的MDC上下文
                MDC.clear();
            }
        };
    }
    
    /**
     * 在当前线程中设置traceId
     * @param traceId traceId
     */
    public static void setTraceId(String traceId) {
        MDC.put("traceId", traceId);
    }
    
    /**
     * 在当前线程中设置sTraceId
     * @param sTraceId
     */
    public static void setSTraceId(String sTraceId) {
        MDC.put("sTraceId", sTraceId);
    }
    
    /**
     * 获取当前线程的traceId
     * @return traceId
     */
    public static String getTraceId() {
        return MDC.get("traceId");
    }
    
    /**
     * 获取当前线程的sTraceId
     * @return sTraceId
     */
    public static String getSTraceId() {
        return MDC.get("sTraceId");
    }
    

    /**
     * 清理当前线程的MDC上下文
     */
    public static void clear() {
        MDC.clear();
    }
    
    /**
     * 包装ExecutorService，使其支持MDC上下文传递
     * @param executorService 原始ExecutorService
     * @return 支持MDC上下文传递的ExecutorService
     */
    public static ExecutorService wrapExecutorService(ExecutorService executorService) {
        return new WrappedExecutorService(executorService);
    }
    
    /**
     * 包装的ExecutorService实现，支持MDC上下文传递
     */
    private static class WrappedExecutorService implements ExecutorService {
        private final ExecutorService executorService;
        
        public WrappedExecutorService(ExecutorService executorService) {
            this.executorService = executorService;
        }
        
        @Override
        public void execute(Runnable command) {
            executorService.execute(MDCUtil.wrap(command));
        }
        
        // 委托给原始ExecutorService的其他方法
        @Override
        public void shutdown() {
            executorService.shutdown();
        }
        
        @Override
        public List<Runnable> shutdownNow() {
            return executorService.shutdownNow();
        }
        
        @Override
        public boolean isShutdown() {
            return executorService.isShutdown();
        }
        
        @Override
        public boolean isTerminated() {
            return executorService.isTerminated();
        }
        
        @Override
        public boolean awaitTermination(long timeout, TimeUnit unit) throws InterruptedException {
            return executorService.awaitTermination(timeout, unit);
        }
        
        @Override
        public <T> Future<T> submit(Callable<T> task) {
            return executorService.submit(MDCUtil.wrap(task));
        }
        
        @Override
        public <T> Future<T> submit(Runnable task, T result) {
            return executorService.submit(MDCUtil.wrap(task), result);
        }
        
        @Override
        public Future< ? > submit(Runnable task) {
            return executorService.submit(MDCUtil.wrap(task));
        }
        
        @Override
        public <T> List<Future<T>> invokeAll(
                Collection<? extends Callable<T>> tasks) throws InterruptedException {
            // 包装所有任务
            List<Callable<T>> wrappedTasks = new ArrayList<>();
            for (Callable<T> task : tasks) {
                wrappedTasks.add(MDCUtil.wrap(task));
            }
            return executorService.invokeAll(wrappedTasks);
        }
        
        @Override
        public <T> List<Future<T>> invokeAll(
                Collection<? extends Callable<T>> tasks, long timeout, TimeUnit unit)
                throws InterruptedException {
            // 包装所有任务
            List<Callable<T>> wrappedTasks = new ArrayList<>();
            for (Callable<T> task : tasks) {
                wrappedTasks.add(MDCUtil.wrap(task));
            }
            return executorService.invokeAll(wrappedTasks, timeout, unit);
        }
        
        @Override
        public <T> T invokeAny(Collection<? extends Callable<T>> tasks)
                throws InterruptedException, ExecutionException {
            // 包装所有任务
            List<Callable<T>> wrappedTasks = new ArrayList<>();
            for (Callable<T> task : tasks) {
                wrappedTasks.add(MDCUtil.wrap(task));
            }
            return executorService.invokeAny(wrappedTasks);
        }
        
        @Override
        public <T> T invokeAny(Collection<? extends Callable<T>> tasks, long timeout,
                               TimeUnit unit)
                throws InterruptedException, ExecutionException, TimeoutException {
            // 包装所有任务
            List<Callable<T>> wrappedTasks = new ArrayList<>();
            for (Callable<T> task : tasks) {
                wrappedTasks.add(MDCUtil.wrap(task));
            }
            return executorService.invokeAny(wrappedTasks, timeout, unit);
        }
    }
}